#include "bench.hpp"
#include <random>
#include <algorithm>

void init_particles(std::vector<ParticleAoS>& aos,
                    ParticlesSoA& soa,
                    std::uint64_t seed,
                    float lo, float hi) {
    std::size_t n = aos.size();
    if (soa.size() != n) {
        soa = ParticlesSoA(n);
    }
    std::mt19937 rng(static_cast<std::mt19937::result_type>(seed));
    std::uniform_real_distribution<float> dist_f(lo, hi);
    std::uniform_real_distribution<double> dist_d(static_cast<double>(lo),
                                                  static_cast<double>(hi));

    for (std::size_t i = 0; i < n; ++i) {
        ParticleAoS p{dist_f(rng), dist_f(rng), dist_f(rng), dist_f(rng), dist_d(rng)};
        aos[i] = p;
        soa.x[i] = p.x;
        soa.y[i] = p.y;
        soa.z[i] = p.z;
        soa.w[i] = p.w;
        soa.mass[i] = p.mass;
    }
}

float bench_aos(std::vector<ParticleAoS>& aos,
                int iters,
                float a, float b, float c, float d,
                double e, double f) {
    float checksum = 0.0f;
    const std::size_t N = aos.size();
    for (int t = 0; t < iters; ++t) {
        for (std::size_t i = 0; i < N; ++i) {
            ParticleAoS& p = aos[i];
            p.x = p.x * a + p.y * b;
            p.y = p.y * b + p.z * c;
            p.z = p.z * c + p.w * d;
            p.w = p.w * d + p.x * a;
            p.mass = p.mass * e + static_cast<double>(p.x) * f; // 使用 mass
            checksum += p.x + p.y + p.z + p.w + static_cast<float>(p.mass);
        }
    }
    return checksum;
}

float bench_soa(ParticlesSoA& soa,
                int iters,
                float a, float b, float c, float d,
                double e, double f) {
    float checksum = 0.0f;
    const std::size_t N = soa.size();
    auto& x = soa.x; auto& y = soa.y; auto& z = soa.z; auto& w = soa.w; auto& m = soa.mass;
    for (int t = 0; t < iters; ++t) {
        for (std::size_t i = 0; i < N; ++i) {
            x[i] = x[i] * a + y[i] * b;
            y[i] = y[i] * b + z[i] * c;
            z[i] = z[i] * c + w[i] * d;
            w[i] = w[i] * d + x[i] * a;
            m[i] = m[i] * e + static_cast<double>(x[i]) * f; // 使用 mass
            checksum += x[i] + y[i] + z[i] + w[i] + static_cast<float>(m[i]);
        }
    }
    return checksum;
}